from datetime import datetime
from contextlib import contextmanager
from sqlalchemy import Column, Integer, SmallInteger
from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy, BaseQuery
from app.libs.error_code import SaveError
from app.libs.error_code import NotFound


class SQLAlchemy(_SQLAlchemy):
    @contextmanager
    def auto_commit(self):
        # __enter__
        try:
            # 自己的逻辑
            yield
            # __exit__
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            # current_app.logger.error(e)
            SaveError()


class Query(BaseQuery):
    def filter_by(self, **kwargs):
        if 'status' not in kwargs.keys():
            kwargs['status'] = 1
        return super(Query, self).filter_by(**kwargs)

    def get_or_404(self, ident):
        rv = self.get(ident)
        if rv is None:
            NotFound(msg='未找到数据')
        return rv

    def first_or_404(self):
        rv = self.first()
        if rv is None:
            NotFound(msg='未找到数据')
        return rv


db = SQLAlchemy(query_class=Query)


class Base(db.Model):
    __abstract__ = True
    create_time = Column('create_time', Integer)
    status = Column(SmallInteger, default=1)

    def __init__(self):
        self.create_time = datetime.now()

    def __getitem__(self, item):
        # getattr()取到当前对象的属性值
        return getattr(self, item) # item就是属性名

    @property
    def create_datetime(self):
        if self.create_time:
            return datetime.fromtimestamp(self.create_time)
        else:
            return None

    def delete(self):
        self.status = 0

    def set_attrs(self, attrs):
        for key, value in attrs.items():
            if hasattr(self, key) and key != 'id':
                setattr(self, key, value)

    def keys(self):
        return self.fields

    def hide(self, *keys):
        for key in keys:
            self.fields.remove(key)  # 将每一个模型中不需要的属性进行删除,最后返回模型
        return self
